import numpy as np
from matplotlib import pyplot as plt


def linear_regression_by_least_squares(x, y):
    """
    线性回归 最小二乘法 (n 个特征数据, 1 个目标值)

    :param x: 特征数据
    :param y: 目标值
    """
    # 归一化
    x, mu, sigma = feature_normalize(x)
    # 处理特征数据与目标值
    x = np.column_stack((np.ones(x.shape[0]), x))  # 特征数据添加截距列 (全1)
    y = y.reshape(-1, 1)  # 将目标值重塑为 shape = (y.shape, 1) 的二维数组 (列向量)
    # 将特征数据与目标值由二维数组转换为矩阵
    x_mat = np.mat(x)
    y_mat = np.mat(y)
    # 求解最小二乘法系数 omega
    omega = (x_mat.T * x_mat).I * x_mat.T * y_mat

    # 绘图查看拟合情况
    ax = plt.axes(projection="3d")
    ax.scatter3D(x_mat[:, 1], x_mat[:, 2], y_mat, color="red")
    ax.scatter3D(x_mat[:, 1], x_mat[:, 2], x_mat * omega, color="green")
    plt.show()

    # 结果预测
    test = np.array([1650, 3]).reshape(1, -1)  # 将测试数据重塑为 shape = (1, test.shape) 的二维数组 (行向量)
    test = (test - mu) / sigma  # 归一化处理测试数据
    test = np.column_stack((np.ones(test.shape[0]), test))  # 测试数据添加截距列 (全1)
    test_mat = np.mat(test)  # 将测试数据转换为矩阵
    print(f"预测结果: {test_mat * omega}")


def feature_normalize(numpy_array):
    """
    归一化处理特征数据

    :param numpy_array: 特征数据 (type: numpy.ndarray)
    :return: 归一化处理后的特征数据、均值向量和标准差向量
    """
    # 计算特征数据每一列的均值和标准差
    mu = np.mean(numpy_array, axis=0)
    sigma = np.std(numpy_array, axis=0)
    # 计算归一化处理后的特征数据
    normalized_array = (numpy_array - mu) / sigma

    return normalized_array, mu, sigma


if __name__ == "__main__":
    data = np.loadtxt("data.txt", delimiter=",")
    x_data = np.array(data[:, 0:-1])
    y_data = np.array(data[:, -1])

    linear_regression_by_least_squares(x_data, y_data)
